from dash import Dash
from dash import dcc, html
from dash.dependencies import Input, Output
import plotly.graph_objects as go
import networkx as nx
import numpy as np
import matplotlib.pyplot as plt
import json
import os

exp_dir = "../../exp/cmpo/discovery-time_scale-1-long-rollout/gym___Hopper-v2/2022.04.30/164209"
mask_dir = os.path.join(exp_dir, "mask")

if not os.path.exists(mask_dir):
    raise "mask dir not exists."


def create_graph(adjacency_matrix):
    input_dim, output_dim = adjacency_matrix.shape
    gr = nx.Graph()
    for i in range(input_dim):
        gr.add_node("i{}".format(i), pos=(1, -i / input_dim))
    for o in range(output_dim):
        gr.add_node("o{}".format(o), pos=(2, -o / output_dim))
    for i, o in zip(*np.where(adjacency_matrix == 1)):
        gr.add_edge("i{}".format(i), "o{}".format(o))

    pos = nx.get_node_attributes(gr, 'pos')

    # edges trace
    edge_x = []
    edge_y = []
    for edge in gr.edges():
        x0, y0 = pos[edge[0]]
        x1, y1 = pos[edge[1]]
        edge_x.append(x0)
        edge_x.append(x1)
        edge_x.append(None)
        edge_y.append(y0)
        edge_y.append(y1)
        edge_y.append(None)

    edge_trace = go.Scatter(
        x=edge_x, y=edge_y,
        line=dict(color='black', width=1),
        hoverinfo='none',
        showlegend=False,
        mode='lines')

    # nodes trace
    node_x = []
    node_y = []
    text = []
    for node in gr.nodes():
        x, y = pos[node]
        node_x.append(x)
        node_y.append(y)
        text.append(node)

    node_trace = go.Scatter(
        x=node_x, y=node_y, text=text,
        mode='markers+text',
        showlegend=False,
        hoverinfo='none',
        marker=dict(
            color='pink',
            size=50,
            line=dict(color='black', width=1)))

    # layout
    layout = dict(plot_bgcolor='white',
                  paper_bgcolor='white',
                  width=1600,
                  height=1000,
                  margin=dict(t=10, b=10, l=10, r=10, pad=0),
                  xaxis=dict(linecolor='black',
                             showgrid=False,
                             showticklabels=False,
                             mirror=True),
                  yaxis=dict(linecolor='black',
                             showgrid=False,
                             showticklabels=False,
                             mirror=True))

    # figure
    fig = go.Figure(data=[edge_trace, node_trace], layout=layout)

    return fig


external_stylesheets = ['https://codepen.io/chriddyp/pen/bWLwgP.css']

app = Dash(__name__, external_stylesheets=external_stylesheets)

all_mask = sorted(os.listdir(mask_dir), key=lambda x: int(x.split(".")[0]))

app.layout = html.Div([
    dcc.Dropdown(
        all_mask,
        all_mask[0],
        id='mask'
    ),
    dcc.Graph(id='graph'),
]
)


@app.callback(
    Output('graph', 'figure'),
    Input('mask', 'value'),
)
def update_output(mask):
    with open(os.path.join(mask_dir, mask)) as f:
        mask_dict = np.array(json.load(f)["causal_mask"])
    return create_graph(mask_dict)


if __name__ == '__main__':
    app.run_server(debug=True)
